#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient


def InsertIdAndName():
    DB_NAME = "mesh"
    COLLECTION_NAME = "descriptor"

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    infile = open("desc.name.total", "r")
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        id = splitline[1]
        name = splitline[2]
        splitid = id.split(".")
        exact_id = splitid[0]
        if len(splitid) == 2:
            insert_content = {}
            insert_content["id"] = exact_id
            insert_content["name"] = name
            insert_content["alias"] = [name]
            # print insert_content
            collection.insert_one(insert_content)
        elif len(splitid) == 3:
            # print exact_id+"\t"+name
            collection.update_one({"id": exact_id}, {"$addToSet": {"alias": name}})
    infile.close()
    connection.close()


def InsertTree():
    DB_NAME = "mesh"
    COLLECTION_NAME = "descriptor"

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]
    infile = open("desc.tree.total", "r")
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        id = splitline[0]
        tree = splitline[1]
        tree_root = tree.split(".")[0]
        collection.update_one({"id": id}, {"$addToSet": {"tree": tree, "tree_root": tree_root}})
        # print id+"\t"+tree_root+"\t"+tree
    infile.close()
    connection.close()


if __name__ == "__main__":
    # InsertIdAndName()
    InsertTree()
